import torch
from itertools import zip_longest


class InputFormatter:

    def __init__(self, max_len, dic, fillvalue):
        self.max_len = max_len
        self.dic = dic
        self.fillvalue = fillvalue

    def trim(self, texts):
        trimmed = []
        for text in texts:
            if len(text) > self.max_len:
                text = text[:self.max_len]
            trimmed.append(text)
        return trimmed

    def indices_from_texts(self, texts):
        """
        :param sentences: list of list，元素为已切分过的句子列表，e.g.:[['营销', '文案', '生成'], ['今天', '天气', '不错', '啊']]
        :return: list of list，元素为已转化为词索引的句子列表，具体的词索引的值取决于字典，e.g.:[[2, 4, 67], [79, 58, 11, 101]]
        """
        return [[self.dic[word] for word in t] for t in texts]

    def add_se_to_texts(self, texts_indices):
        return [[self.dic.SOS] + t + [self.dic.EOS] for t in texts_indices]

    def zero_padding(self, indices_list, fillvalue=0):
        """
        :param indices_list:  list of list，元素为已转化为词索引的句子列表，具体的词索引的值取决于字典，e.g.:[[2, 4, 67], [[79, 58, 11, 101]]]
        :param fillvalue: 使用字典的PAD索引进行长度填充，长度由最长的句子决定
        :return:
        """
        padded = list(zip_longest(*indices_list, fillvalue=fillvalue))
        return torch.tensor(padded).transpose(1, 0)

    def flatten_2d(self, texts):
        flattened = []
        for text in texts:
            flattened.extend(text)
        return flattened

    def __call__(self, texts, fold=None):
        if fold:
            self.max_len *= fold
            indices = torch.arange(len(texts))
            starts = torch.tensor(indices % fold, dtype=torch.int)
            start_indices = torch.where(starts == 0)[0]
            start_indices = list(map(lambda idx: int(idx), start_indices))

            texts = list(map(lambda s: self.flatten_2d(texts[s:s+fold]), start_indices))
            print("texts:", texts[:5])
        trimmed = self.trim(texts)
        indice_list = self.indices_from_texts(trimmed)
        indice_list = self.add_se_to_texts(indice_list)
        padded = self.zero_padding(indice_list)
        return padded


